package com.zst.common.utils;

import sun.misc.BASE64Decoder;
import sun.misc.BASE64Encoder;

import javax.crypto.*;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;

/**
 * @author zst
 */
public class AESUtils {
    private static final String RANDOM_ALGORITHM = "SHA1PRNG";
    private static final String KEY_ALGORITHM = "AES";
    private static final int KEY_LENGTH = 128;
    private static final String CIPHER_ALGORITHM = "AES/ECB/PKCS5Padding";

    public static String encrypt(String key, String content) throws Exception {
        if (key == null || key.length() == 0) {
            throw new IllegalArgumentException("Key cannot null");
        }
        if (content == null || content.length() == 0) {
            throw new IllegalArgumentException("Content cannot null");
        }
        String ret = "";
        try {
            // 创建AES密钥
            Key k = buildAESKey(key);
            Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, k);
            byte[] result = cipher.doFinal(content.getBytes(StandardCharsets.UTF_8));
            BASE64Encoder encoder = new BASE64Encoder();
            ret = encoder.encode(result);
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | BadPaddingException | IllegalBlockSizeException | InvalidKeyException e) {
            throw e;
        }
        return ret;
    }

    public static String decrypt(String key, String content) throws Exception {
        if (key == null || key.length() == 0) {
            throw new IllegalArgumentException("Key cannot null");
        }
        if (content == null || content.length() == 0) {
            throw new IllegalArgumentException("Content cannot null");
        }

        String ret = "";
        try {
            Key k = buildAESKey(key);
            Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, k);
            BASE64Decoder decoder = new BASE64Decoder();
            byte[] decryptResult = cipher.doFinal(decoder.decodeBuffer(content));
            ret = new String(decryptResult, StandardCharsets.UTF_8);
        } catch (NoSuchAlgorithmException | IOException | BadPaddingException | IllegalBlockSizeException | InvalidKeyException | NoSuchPaddingException e) {
            throw e;
        }

        return ret;
    }

    /**
     * 将字符串转换为AES算法使用的key
     * @param originKey 原始key字符串
     * @throws NoSuchAlgorithmException
     */
    public static Key buildAESKey(String originKey) throws NoSuchAlgorithmException {
        try {
            SecureRandom secureRandom = SecureRandom.getInstance(RANDOM_ALGORITHM);
            secureRandom.setSeed(originKey.getBytes(StandardCharsets.UTF_8));
            KeyGenerator generator = KeyGenerator.getInstance(KEY_ALGORITHM);
            generator.init(KEY_LENGTH, secureRandom);
            return generator.generateKey();
        } catch (NoSuchAlgorithmException e) {
            throw e;
        }
    }
}
